#!/usr/bin/env python3

import pymongo
import numpy as np
import json
from pprint import pprint

client = pymongo.MongoClient('192.168.0.12', 27117)

db = client.hntp

col1 = db.tp_2km_grids
col2 = db.tp_2km_zlev
col3 = db.tp_2km_yearly

grids = col1.find({
	'location': {
		'$geoWithin': {
			'$geometry': {
				'type': 'Polygon',
				'coordinates': [
					[
						[101, 29],
						[101.5, 29],
						[101.5, 29.5],
						[101, 29.5],
						[101, 29]
					]
				]
			}
		}
	}
})

grid_ids = [x['_id'] for x in grids]

zlev = col2.find_one()['zlev']

data = list(col3.find({
	'grid_id': { '$in': grid_ids }
}))

band_wdir_freq = (0, 22.5, 45, 67.5, 90, 112.5, 135, 157.5, 180, 202.5, 225, 247.5, 270, 292.5, 315, 337.5)
band_wsp_freq = (1.5, 3.5, 4.5, 5.5, 6.5, 7.5, 8.5, 9.5, 10.5, 11.5, 12.5, 13.5, 14.5, 15.5, 16.5, 17.5)

res = {}
fields = {
	't': 1,
	'pres': 1,
	'q': 1,
	'wsp': 1,
	'wdir_freq': len(band_wdir_freq),
	'wsp_freq': len(band_wsp_freq)
}
for field in fields:
	res[field] = {}
for k, z in enumerate(zlev):
	for field, size in fields.items():
		res[field][z] = np.zeros(size)
	for x in data:
		for field in fields:
			res[field][z] += np.array(x[field][k])
	for field, size in fields.items():
		if size == 1:
			res[field][z] = float(res[field][z]) / len(data)
		else:
			res[field][z] /= len(data)
res['time'] = data[0]['time'].isoformat()

tmp = {}
for z, x in res['wdir_freq'].items():
	tmp[z] = {}
	for b, y in enumerate(x):
		tmp[z][band_wdir_freq[b]] = y
res['wdir_freq'] = tmp
tmp = {}
for z, x in res['wsp_freq'].items():
	tmp[z] = {}
	for b, y in enumerate(x):
		tmp[z][band_wsp_freq[b]] = y
res['wsp_freq'] = tmp

open('test.json', 'w').write(json.dumps(res))
